# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

version: '3'

services:
  sdr:
    container_name: fm-rag-sdr
    image: fm-rag-sdr-holoscan:latest
    build:
      context: ${PROJECT_DIR?:source compose.env}/sdr-holoscan
      dockerfile: Dockerfile

    restart: "no"

    environment:
      TZ: ${TIMEZONE:-America/New_York}
      SDR_LOG_LEVEL: ${SDR_LOG_LEVEL:-WARN}
      FRONTEND_URI: ${FRONTEND_URI}
      DATABASE_URI: ${CHAIN_SERVER_HOST}:${CHAIN_SERVER_PORT}

    working_dir: /workspace/
    command: "python sdr-holoscan/app.py"

    network_mode: host
    devices:
      - "/dev/bus/usb:/dev/bus/usb"
      - "/dev/snd:/dev/snd"

    # Enable GPU usage
    runtime: nvidia
    shm_size: 8gb
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['${SDR_GPU:-all}']
              capabilities: [gpu]

  frontend:
    container_name: fm-rag-frontend
    image: fm-rag-frontend:latest
    build:
      context: ${PROJECT_DIR?:source compose.env}/frontend
      dockerfile: Dockerfile

    command: --port ${FRONTEND_SERVER_PORT}

    environment:
      TZ: ${TIMEZONE:-America/New_York}
      FRONTEND_LOG_LEVEL: ${FRONTEND_LOG_LEVEL:-WARN}
      FRONTEND_URI: ${FRONTEND_URI}
      APP_SERVERURL: http://localhost
      APP_SERVERPORT: ${CHAIN_SERVER_PORT}

    ports:
      - "${CHAIN_SERVER_PORT}:${CHAIN_SERVER_PORT}"
    expose:
      - "${CHAIN_SERVER_PORT}"
    network_mode: host

    deploy:
      resources:
        reservations:
          devices:
          - driver: nvidia
            device_ids: ['${FRONTEND_GPU:-all}']
            capabilities: [gpu]

  server:
    container_name: fm-rag-chain-server
    image: fm-rag-chain-server:latest
    build:
      context: ${PROJECT_DIR?:source compose.env}/chain-server
      dockerfile: Dockerfile

    command: --host ${CHAIN_SERVER_HOST} --port ${CHAIN_SERVER_PORT}

    environment:
      TZ: ${TIMEZONE:-America/New_York}
      NVIDIA_API_KEY: ${NVIDIA_API_KEY:-}
      CHAIN_LOG_LEVEL: ${CHAIN_LOG_LEVEL:-WARN}
      DB_CHUNK_SIZE: ${DB_CHUNK_SIZE:-256}
      DB_CHUNK_OVERLAP: ${DB_CHUNK_OVERLAP:-32}

    ports:
      - "8081:8081"
    expose:
      - "8081"
    network_mode: host

    deploy:
      resources:
        reservations:
          devices:
          - driver: nvidia
            device_ids: ['${CHAIN_GPU:-all}']
            capabilities: [gpu]

  replay:
    container_name: fm-rag-file-replay
    image: fm-rag-file-replay:latest
    build:
      context: ${PROJECT_DIR?:source compose.env}/file-replay
      dockerfile: Dockerfile

    environment:
      TZ: ${TIMEZONE:-America/New_York}

    volumes:
      - ${PROJECT_DIR}/file-replay/files:/workspace/files

    working_dir: /workspace/
    command: >
      python wav_replay.py
        --file-name ${REPLAY_FILE}
        --dst-ip ${SDR_IP:-"0.0.0.0"}
        --dst-port ${SDR_PORT:-5005}
        --sample-rate ${SDR_SAMPLE_RATE:-1000000}
        --packet-size ${SDR_MAX_PKT_SZ:-1472}
        --total-time ${REPLAY_TIME:-0}

    network_mode: host
    devices:
      - "/dev/bus/usb:/dev/bus/usb"
      - "/dev/snd:/dev/snd"

    # Enable GPU usage
    runtime: nvidia
    shm_size: 8gb
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['${REPLAY_GPU:-all}']
              capabilities: [gpu]

  nim:
    container_name: fm-rag-nim
    image: nim:latest
    build:
      context: ${PROJECT_DIR?:source compose.env}/nim
      dockerfile: Dockerfile

    volumes:
      - ${MODEL_CHECKPOINT}:/huggingface-dir
      - ${NIM_MODEL_PATH}:/model-store
      - ${NIM_CONFIG_FILE}:/model_config.yaml

    environment:
      NIM_OPENAI_PORT: ${NIM_OPENAI_PORT}

    ports:
      - "${NIM_OPENAI_PORT}:${NIM_OPENAI_PORT}"
    expose:
      - "${NIM_OPENAI_PORT}"

    # Start inference server
    command: >
      nemollm_inference_ms
        --model mistral_7b
        --openai_port=${NIM_OPENAI_PORT}
        --num_gpus=1

    # Enable GPU usage
    runtime: nvidia
    shm_size: 8gb
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['${NIM_GPU:-all}']
              capabilities: [gpu]